#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjianrong
@software: pycharm
@file: prune_resnet50.py
@time: 2020/10/19 16:05
@desc:
'''

from torchvision.models.resnet import resnet50
import torch
from torchvision.datasets import MNIST
from torchvision import transforms
from torch.utils.data import DataLoader
from torch_op.common import get_sync_time
import torch.nn.utils.prune as prune



test_dataset = MNIST('torch_op/mnist/data',train=False,download=True,
                      transform=transforms.Compose([
                          transforms.ToTensor(),
                          # transforms.Normalize((0.1307,),(0.3081))
                      ])
                      )

test_loader = DataLoader(test_dataset,batch_size=64,shuffle=False,num_workers=4)

def test():
    model.eval()
    total_right = 0
    total_wrong = 0
    with torch.no_grad():
        for batch_idx,(data,target) in enumerate(test_loader):
            if data.size(1) == 1:
                data = torch.cat([data, data, data], dim=1)
            data = data.cuda()
            target = target.cuda()
            output = model(data)
            output = output.max(dim=1)[1]
            right = torch.sum(output==target).item()
            wrong = len(data) - right
            total_right += right
            total_wrong += wrong
    acc = total_right/(total_wrong+total_right)
    print('acc:{}'.format(acc))

model = resnet50(False,num_classes=10)
model.load_state_dict(torch.load('torch_op/mnist/models/best.pt')['net'])
model.cuda()

s = get_sync_time()
test()
e = get_sync_time()
print(e-s)

parameters_to_prune = []
for module in model.modules():
    parameters_to_prune.append(module)
parameters_to_prune = tuple(parameters_to_prune)

# parameters_to_prune = (
#     (model.conv1, 'weight'),
#     (model.conv2, 'weight'),
#     (model.fc1, 'weight'),
#     (model.fc2, 'weight'),
#     (model.fc3, 'weight'),
# )
#
# prune.global_unstructured(
#     parameters_to_prune,
#     pruning_method=prune.L1Unstructured,
#     amount=0.9,
# )
#
#
#


